import numpy as np
import matplotlib.pyplot as plt
import math

# 画图设置
plt.style.use('default')
plt.rcParams['font.family'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 定义一个函数，用于生成正弦波
def getSin(amp, freq, phase, sampleList):
    return amp * np.sin(2 * math.pi * freq * sampleList + phase)

# 定义一个函数，用于生成余弦波
def getCos(amp, freq, phase, sampleList):
    return amp * np.cos(2 * math.pi * freq * sampleList + phase)

# 1. 获得混合波形
srate = 3000           # 设置采样率为 3000
t = np.linspace(0, 1, srate)     # 在 0 到 1 的范围内生成等间距的采样点

# 分别生成正、余弦波
s1 = getSin(amp=1.5, freq=-30, phase=0, sampleList=t)
s2 = getCos(amp=3, freq=-5, phase=0, sampleList=t)
s3 = getSin(amp=10, freq=-100, phase=0, sampleList=t)
s4 = getCos(amp=20, freq=-120, phase=0, sampleList=t)
m= s1+s2+s3+s4             # 将正弦波相加，得到混合波形

# 2.  通过fft获得混合波形的频域图

# 获得傅里叶系数
fCoefs = np.fft.fft(m,srate)

# 获得振幅列表：每一个绕线的重心到原点的距离
amp_list=2*np.abs(fCoefs/srate)

# 把频率轴从0~3000 转变成 0~1500 然后 -1500~-1
freqs = np.fft.fftfreq(len(amp_list), 1/srate)

# 然后把 频率轴 和 数据 都变成 0hz 在中间，向左是负频率，向右是正频率的形式
amp_shifted=np.fft.fftshift(amp_list)
freq_shift=np.fft.fftshift(freqs)

# 通过mask得到过滤掉噪音的频域图
amp_list[:-110][110:]=0
mask = amp_list > 0.5
fCoef_clean = mask * fCoefs
B_clean = np.fft.ifft(fCoef_clean)

#双重去噪
amp_list1=np.copy(amp_list)
mask = (amp_list1>1)
list = mask*amp_list1        #大于1的归0
list[:-110][110:]=0         #选取大于110hz为0

# 然后把 频率轴 和 数据 都变成 0hz 在中间，向左是负频率，向右是正频率的形式
amp_shifted1=np.fft.fftshift(list)
freq_shift1=np.fft.fftshift(freqs)

fg,ax=plt.subplots(2,2,figsize=(20,5))

# 绘制噪音混合波形m时域图选取[0~500]的数据显示
ax[0,0].grid()
ax[0,0].plot(np.real(m[0:500]))               #np.fft.ifft反馈的是复数结果，想要获得时域图需要利用np.real抽取其中的真值
ax[0,0].set_xlim(0,500)
ax[0,0].set_title("原始波形")

# 绘制去噪混合波形m时域图选取[0~500]的数据显示
ax[0,1].grid()
ax[0,1].plot(np.real(B_clean[0:500]))
ax[0,1].set_xlim(0,500)
ax[0,1].set_title("去噪后的波形")

ax[1,0].grid()
ax[1,0].stem(freq_shift,amp_shifted)
ax[1,0].set_xlim([-150,150])
ax[1,0].set_title("原始振幅谱")

ax[1,1].grid()
ax[1,1].stem(freq_shift1,amp_shifted1)
ax[1,1].set_xlim([-150,150])
ax[1,1].sharey(ax[1,0])
ax[1,1].set_title("去噪后的振幅谱")
plt.tight_layout()
plt.show()